import numpy as np
from matplotlib import pyplot as plt
# Note: Ensure 'pillow' package is installed for handling image loading
from sklearn.datasets import load_sample_image
from sklearn.cluster import MiniBatchKMeans

# Load the sample image ("china.jpg") from sklearn's datasets
china = load_sample_image("china.jpg")

# Display the original image
ax = plt.axes(xticks=[], yticks=[])  # Remove axis ticks for cleaner visualization
ax.imshow(china)
plt.title("Original Image")
plt.show()

# Inspect the shape of the image
print("Image shape:", china.shape)  # Shape: (427, 640, 3), representing height, width, and RGB channels

# Normalize image data to a scale of 0 to 1
data = china / 255.0
# Reshape the image data from (427, 640, 3) to (427*640, 3) for clustering
data = data.reshape(-1, 3)  # Automatically calculates the first dimension
print("Reshaped data shape:", data.shape)  # Shape: (273280, 3)


# Define a function to visualize pixel colors in RGB space
def plot_pixels(data, title, colors=None, N=10000):
    """
    Plot RGB color space with a random subset of pixels.

    Parameters:
        data (ndarray): The image data in RGB format, shape (N, 3).
        title (str): Title of the plot.
        colors (ndarray, optional): Colors for plotting points. Defaults to None (uses data itself).
        N (int, optional): Number of random pixels to sample. Defaults to 10,000.
    """
    if colors is None:
        colors = data

    # Choose a random subset of pixels
    rng = np.random.RandomState(0)
    i = rng.permutation(data.shape[0])[:N]
    colors = colors[i]
    R, G, B = data[i].T  # Transpose to extract R, G, B channels separately

    # Create scatter plots for color distributions
    fig, ax = plt.subplots(1, 2, figsize=(16, 6))
    ax[0].scatter(R, G, color=colors, marker='.')
    ax[0].set(xlabel='Red', ylabel='Green', xlim=(0, 1), ylim=(0, 1))
    ax[1].scatter(R, B, color=colors, marker='.')
    ax[1].set(xlabel='Red', ylabel='Blue', xlim=(0, 1), ylim=(0, 1))

    fig.suptitle(title, size=20)
    plt.show()


# Visualize the original image's color space
plot_pixels(data, title="Input color space: 16 million possible colors")

# Perform color quantization using K-Means clustering
kmeans = MiniBatchKMeans(n_clusters=16, n_init='auto', random_state=42)
kmeans.fit(data)

# Map each pixel to its closest cluster center
new_colors = kmeans.cluster_centers_[kmeans.predict(data)]

# Visualize the reduced color space
plot_pixels(data, colors=new_colors, title="Reduced color space: 16 colors")

# Reshape the clustered colors back to the original image shape
china_recolored = new_colors.reshape(china.shape)

# Display the original and recolored images side by side
fig, ax = plt.subplots(1, 2, figsize=(16, 6), subplot_kw=dict(xticks=[], yticks=[]))
fig.subplots_adjust(wspace=0.05)  # Reduce space between subplots
ax[0].imshow(china)
ax[0].set_title('Original Image', size=16)
ax[1].imshow(china_recolored)
ax[1].set_title('16-color Image', size=16)
plt.show()